#ifndef MYAMG_CSRSMOOTHERKERNELS_CUH
#define MYAMG_CSRSMOOTHERKERNELS_CUH

namespace myamg {
    template<typename Integer, typename Real>
    void calcRecipDHost(Real *recipD, const Integer *Ap, const Integer *Ai, const Real *Av, const size_t nRows) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; ++i) {
            Real diag = 0;
            for (auto j = Ap[i]; j < Ap[i + 1]; ++j) {
                if (Ai[j] == i) {
                    diag = Av[j];
                    break;
                }
            }
            recipD[i] = (diag != 0) ? (1.0 / diag) : 0.0;
        }
    }

    template<typename Integer, typename Real>
    __global__ void calcRecipDDevice(Real *recipD, const Integer *Ap, const Integer *Ai, const Real *Av,
                                     const size_t nRows) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            Real diag = 0;
            for (auto j = Ap[idx]; j < Ap[idx + 1]; ++j) {
                if (Ai[j] == idx) {
                    diag = Av[j];
                    break;
                }
            }
            recipD[idx] = (diag != 0) ? (1.0 / diag) : 0.0;
        }
    }

    template<typename Integer, typename Real>
    void jacobiUpdateXandZHost(Real *x, Real *z, const Real *recipD, const Integer *Ap, const Integer *Ai,
                               const Real *Av, const Real *b, const size_t nRows, const double omega) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; ++i) {
            auto temp = b[i];
            for (auto j = Ap[i]; j < Ap[i + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            z[i] = temp;
        }
#pragma omp parallel for
        for (auto i = 0; i < nRows; ++i) {
            x[i] += omega * recipD[i] * z[i];
        }
    }

    template<typename Integer, typename Real>
    __global__ void jacobiUpdateZDevice(Real *z, const Real *x, const Integer *Ap, const Integer *Ai,
                                        const Real *Av, const Real *b, const size_t nRows) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            auto temp = b[idx];
            for (auto j = Ap[idx]; j < Ap[idx + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            z[idx] = temp;
        }
    }

    template<typename Real>
    __global__ void jacobiUpdateXDevice(Real *x, const Real *z, const Real *recipD, const size_t nRows,
                                        const double omega) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            x[idx] += omega * recipD[idx] * z[idx];
        }
    }

    template<typename Real>
    void jacobiUpdateXZeroHost(Real *x, const Real *recipD, const Real *b, const size_t nRows, const double omega) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; ++i) {
            x[i] = omega * recipD[i] * b[i];
        }
    }

    template<typename Real>
    __global__ void jacobiUpdateXZeroDevice(Real *x, const Real *recipD, const Real *b, const size_t nRows,
                                            const double omega) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            x[idx] = omega * recipD[idx] * b[idx];
        }
    }

    template<typename Integer, typename Real>
    void calcResidHost(Real *r, const Integer *Ap, const Integer *Ai, const Real *Av,
                       const Real *x, const Real *b, const size_t nRows) {
#pragma omp parallel for
        for (auto i = 0; i < nRows; ++i) {
            auto temp = b[i];
            for (auto j = Ap[i]; j < Ap[i + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            r[i] = temp;
        }
    }

    template<typename Integer, typename Real>
    __global__ void calcResidDevice(Real *r, const Integer *Ap, const Integer *Ai, const Real *Av,
                                    const Real *x, const Real *b, const size_t nRows) {
        const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < nRows) {
            auto temp = b[idx];
            for (auto j = Ap[idx]; j < Ap[idx + 1]; ++j) {
                temp -= Av[j] * x[Ai[j]];
            }
            r[idx] = temp;
        }
    }
}

#endif //MYAMG_CSRSMOOTHERKERNELS_CUH
